__global__
void kernel_gemm_impl1(float const * const __restrict lhs, float const * const __restrict rhs, float * const __restrict out,
                 ::std::uint32_t const M, ::std::uint32_t const K, ::std::uint32_t const N,
                 bool const is_lhs_transpose,
                 bool const is_rhs_transpose,
                 bool const is_out_transpose)
{
    ::std::uint32_t const thread_idx{blockIdx.x * blockDim.x + threadIdx.x};
    if (thread_idx >= N) {
        return;
    }

    tops_dte_ctx_t dte;

    // 初始化DTE
    // FIXME use dte_scope instead
    dte.init();

    for (::std::uint32_t n{thread_idx}; n < N; n += 24)
    {
        for (::std::uint32_t m{}; m < M; ++m)
        {
            float output{};

            for (::std::uint32_t k{}; k < K; ++k)
            {
                float l1_input_l [[clang::uninitialized]];
                ::tops::mdspan l1_input_l_view(tops::Private, __builtin_addressof(l1_input_l), 1);
                if (is_lhs_transpose) {
                    ::tops::mdspan l3_input_l_view(tops::Global, const_cast<float*>(lhs + k * M + m), 1);
                    ::tops::memcpy(dte, l1_input_l_view, l3_input_l_view);
                } else {
                    ::tops::mdspan l3_input_l_view(tops::Global, const_cast<float*>(lhs + m * K + k), 1);
                    ::tops::memcpy(dte, l1_input_l_view, l3_input_l_view);
                }

                float l1_input_r [[clang::uninitialized]];
                ::tops::mdspan l1_input_r_view(tops::Private, __builtin_addressof(l1_input_r), 1);
                if (is_rhs_transpose) {
                    ::tops::mdspan l3_input_r_view(tops::Global, const_cast<float*>(rhs + n * K + k), 1);
                    ::tops::memcpy(dte, l1_input_r_view, l3_input_r_view);
                } else {
                    ::tops::mdspan l3_input_r_view(tops::Global, const_cast<float*>(rhs + k * N + n), 1);
                    ::tops::memcpy(dte, l1_input_r_view, l3_input_r_view);
                }

                output += l1_input_l * l1_input_r;
            }

            ::tops::mdspan l1_output_view(tops::Private, __builtin_addressof(output), 1);
            if (is_out_transpose) {
                ::tops::mdspan l3_output_view(tops::Global, out + n * M + m, 1);
                ::tops::memcpy(dte, l3_output_view, l1_output_view);
            } else {
                ::tops::mdspan l3_output_view(tops::Global, out + m * N + n, 1);
                ::tops::memcpy(dte, l3_output_view, l1_output_view);
            }
        }
    }
}

void GCU_GEMM(float * __restrict dev_lhs,
              float * __restrict dev_rhs,
              float * __restrict dev_out,
              int m, int k, int n,
              bool is_lhs_transpose,
              bool is_rhs_transpose,
              bool is_out_transpose)
{
    if (n <= 12) {
        ::kernel_gemm_impl1<<<1, n>>>(dev_lhs, dev_rhs, dev_out,
            ::std::uint32_t(m), ::std::uint32_t(k), ::std::uint32_t(n),
            is_lhs_transpose,
            is_rhs_transpose,
            is_out_transpose);
    }
    else {
        ::kernel_gemm_impl1<<<2, 12>>>(dev_lhs, dev_rhs, dev_out,
            ::std::uint32_t(m), ::std::uint32_t(k), ::std::uint32_t(n),
                            is_lhs_transpose,
                            is_rhs_transpose,
                            is_out_transpose);
    }
}